import cv2
import numpy as np
import os
import matplotlib.pyplot as plt


def detect_frog_by_color(img_path, visualize=False):
    """
    基于颜色特征检测青蛙并返回ROI区域 - 优化版，特别针对小青蛙
    :param img_path: 图像文件路径
    :param visualize: 是否可视化处理过程
    :return: ROI区域图像和检测状态
    """
    # 读取图像
    img = cv2.imread(img_path)
    if img is None:
        print(f"无法读取图像: {img_path}")
        return None, False

    # 图像增强 - 提高对比度
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    l = clahe.apply(l)
    lab = cv2.merge((l, a, b))
    img_enhanced = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    # 转换为HSV颜色空间
    hsv = cv2.cvtColor(img_enhanced, cv2.COLOR_BGR2HSV)

    # 定义青蛙颜色的HSV范围
    # 1. 标准绿色范围
    lower_green = np.array([25, 30, 30])
    upper_green = np.array([95, 255, 255])

    # 2. 暗黄色范围
    lower_dark_yellow = np.array([10, 40, 20])
    upper_dark_yellow = np.array([40, 255, 180])

    # 3. 浅绿色范围（针对小青蛙）
    lower_light_green = np.array([25, 20, 20])
    upper_light_green = np.array([95, 100, 100])

    # 4. 浅黄色范围（针对小青蛙）
    lower_light_yellow = np.array([10, 20, 20])
    upper_light_yellow = np.array([40, 100, 100])

    # 创建颜色掩膜
    mask_green = cv2.inRange(hsv, lower_green, upper_green)
    mask_dark_yellow = cv2.inRange(hsv, lower_dark_yellow, upper_dark_yellow)
    mask_light_green = cv2.inRange(hsv, lower_light_green, upper_light_green)
    mask_light_yellow = cv2.inRange(hsv, lower_light_yellow, upper_light_yellow)

    # 合并所有掩膜
    combined_mask = cv2.bitwise_or(mask_green, mask_dark_yellow)
    combined_mask = cv2.bitwise_or(combined_mask, mask_light_green)
    combined_mask = cv2.bitwise_or(combined_mask, mask_light_yellow)

    # 形态学操作
    kernel = np.ones((3, 3), np.uint8)  # 更小的内核
    cleaned_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel)
    cleaned_mask = cv2.morphologyEx(cleaned_mask, cv2.MORPH_CLOSE, kernel)

    # 查找轮廓
    contours, _ = cv2.findContours(cleaned_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 如果没有找到轮廓，返回失败
    if not contours:
        return None, False

    # 筛选轮廓
    min_contour_area = 300  # 降低面积阈值
    valid_contours = []

    for contour in contours:
        area = cv2.contourArea(contour)
        if area < min_contour_area:
            continue

        # 计算轮廓的边界矩形
        x, y, w, h = cv2.boundingRect(contour)

        # 计算宽高比
        aspect_ratio = w / float(h)

        # 过滤掉太细长或太扁平的轮廓
        if aspect_ratio < 0.2 or aspect_ratio > 5.0:
            continue

        valid_contours.append(contour)

    # 如果没有有效轮廓，返回失败
    if not valid_contours:
        return None, False

    # 选择面积最大的有效轮廓
    largest_contour = max(valid_contours, key=cv2.contourArea)

    # 获取边界框
    x, y, w, h = cv2.boundingRect(largest_contour)

    # 确保边界框合理大小（至少占图像的3%，降低阈值）
    min_size = 0.03 * img.shape[0] * img.shape[1]
    if w * h < min_size:
        return None, False

    # 扩展边界框（增加15%的边距）
    margin = 0.15
    x = max(0, int(x - margin * w))
    y = max(0, int(y - margin * h))
    w = min(img.shape[1] - x, int(w * (1 + 2 * margin)))
    h = min(img.shape[0] - y, int(h * (1 + 2 * margin)))

    # 提取ROI
    roi = img[y:y + h, x:x + w]

    # 可视化处理过程
    if visualize:
        # 创建可视化图像
        vis_img = img.copy()
        cv2.rectangle(vis_img, (x, y), (x + w, y + h), (0, 255, 0), 3)

        # 创建彩色掩膜显示
        color_mask = np.zeros_like(img)
        color_mask[mask_green > 0] = [0, 255, 0]  # 绿色区域
        color_mask[mask_dark_yellow > 0] = [0, 165, 255]  # 橙色表示暗黄色区域
        color_mask[mask_light_green > 0] = [0, 100, 0]  # 深绿色表示浅绿色区域
        color_mask[mask_light_yellow > 0] = [0, 100, 255]  # 浅蓝色表示浅黄色区域
        cleaned_mask_3d = cv2.cvtColor(cleaned_mask, cv2.COLOR_GRAY2BGR)
        # 组合显示
        top_row = np.hstack((img, cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)))
        middle_row = np.hstack((color_mask, cleaned_mask_3d))
        roi_resized = cv2.resize(roi, (vis_img.shape[1], vis_img.shape[0]))  # 调整 roi 的宽度和 vis_img 高度一致
        bottom_row = np.hstack((vis_img, roi_resized))
        combined_vis = np.vstack((top_row, middle_row, bottom_row))

        # 调整大小以适应屏幕
        scale = min(1.0, 1000 / combined_vis.shape[1])
        combined_vis = cv2.resize(combined_vis, (0, 0), fx=scale, fy=scale)

        # 保存可视化结果
        output_dir = os.path.dirname(img_path) + "_visualization"
        os.makedirs(output_dir, exist_ok=True)
        vis_path = os.path.join(output_dir, f"vis_{os.path.basename(img_path)}")
        cv2.imwrite(vis_path, combined_vis)
        print(f"已保存可视化结果: {vis_path}")

    return roi, True


def preprocess_images_with_color(image_dir, output_dir):
    """
    使用颜色分割预处理图像：只保存成功检测到青蛙的ROI
    """
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 创建日志文件
    log_file = os.path.join(output_dir, "processing_log.txt")
    with open(log_file, "w") as f:
        f.write("图像处理日志\n")
        f.write("=" * 50 + "\n")

    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))]
    saved_count = 0

    print(f"开始处理 {len(image_files)} 张图像...")

    for img_file in image_files:
        img_path = os.path.join(image_dir, img_file)
        roi, detected = detect_frog_by_color(img_path, visualize=True)

        if detected and roi is not None:
            roi_path = os.path.join(output_dir, img_file)
            cv2.imwrite(roi_path, roi)
            saved_count += 1
            print(f"已保存ROI: {roi_path}")
            with open(log_file, "a") as f:
                f.write(f"成功: {img_file}\n")
        else:
            print(f"未检测到青蛙: {img_file}")
            with open(log_file, "a") as f:
                f.write(f"失败: {img_file}\n")

    success_rate = saved_count / len(image_files) * 100

    print(f"\n处理完成: 共处理 {len(image_files)} 张图片")
    print(f"成功检测并保存 {saved_count} 张青蛙ROI图像")
    print(f"未能检测到青蛙 {len(image_files) - saved_count} 张")
    print(f"成功率: {success_rate:.2f}%")

    with open(log_file, "a") as f:
        f.write("\n" + "=" * 50 + "\n")
        f.write(f"总图像数: {len(image_files)}\n")
        f.write(f"成功检测数: {saved_count}\n")
        f.write(f"失败检测数: {len(image_files) - saved_count}\n")
        f.write(f"成功率: {success_rate:.2f}%\n")

    return success_rate


if __name__ == '__main__':
    image_directory = 'D:/Frog/image4'
    output_directory = 'D:/Frog/SmallFrog2'

    # 处理图像
    success_rate = preprocess_images_with_color(image_directory, output_directory)
    print(f"颜色分割检测完成！成功率: {success_rate:.2f}%")